#!/usr/bin/env python
# encoding: utf-8
"""
@author: zhifeng
@license: (C) Copyright 2013-2017, Vivo Corporation Limited.
@contact: zhifeng.ding@vivo.com
@software: image vision
@file: literassp.py
@time: 2020/7/15 17:36
@desc:
"""
import tensorflow as tf
from .network import ConvBlock

class LiteRASSP(tf.keras.layers.Layer):
    def __init__(self, shape=(224, 224), num_classes=4, avg_pool_kernel=(11, 11), avg_pool_strides=(4, 4),
                 resize_method=tf.image.ResizeMethod.BILINEAR, **kwargs):
        """LiteRASSP.
        # Arguments
            # init
                input_shape: Tuple/list of 2 integers, spatial shape of input tensor
                n_class: Integer, number of classes.
                avg_pool_kernel: Tuple/integer, size of the kernel for AveragePooling
                avg_pool_strides: Tuple/integer, stride for applying the of AveragePooling operation
            # Call
                inputs: Tensor, input tensor of the model
                training: Mode for training-aware layers
        # Returns
            Output tensor of the original shape
            """
        super(LiteRASSP, self).__init__(name="LiteRASSP")
        self.shape = shape
        self.n_class = num_classes
        self.avg_pool_kernel = avg_pool_kernel#11
        self.avg_pool_strides = avg_pool_strides#4
        self.resize_method = resize_method
        #branch1
        self.branch1_convblock = ConvBlock(128, 1, strides=1, activation=tf.nn.relu)
        #branch2
        self.branch2_avgpool = tf.keras.layers.AveragePooling2D(pool_size=self.avg_pool_kernel,
                                                strides=self.avg_pool_strides)
        self.branch2_conv = tf.keras.layers.Conv2D(128, 1, strides=1)
        #bracnh3
        self.branch3_conv = tf.keras.layers.Conv2D(self.n_class, 1, strides=1)
        #merge1_2
        self.merge1_2_conv = tf.keras.layers.Conv2D(self.n_class, 1, strides=1)

    def call(self, inputs, training=True):
        out_feature8, out_feature16 = inputs
        # branch1
        x1 = self.branch1_convblock(out_feature16, training=training)
        # branch2
        s = x1.shape
        x2 = self.branch2_avgpool(out_feature16)
        x2 = self.branch2_conv(x2)
        x2 = tf.keras.layers.Activation('sigmoid')(x2)
        x2 = tf.image.resize(x2,
                             size=(int(s[1]), int(s[2])),
                             method=self.resize_method,
                             preserve_aspect_ratio=False,
                             antialias=False,
                             name=None)
        # branch3
        x3 = self.branch3_conv(out_feature8)
        # merge1_2
        x = tf.keras.layers.Multiply()([x1, x2])
        x = tf.image.resize(x,
                            size=(int(2*s[1]), int(2*s[2])),
                            method=self.resize_method,
                            preserve_aspect_ratio=False,
                            antialias=False,
                            name=None)
        x = self.merge1_2_conv(x)
        # merge3
        x = tf.keras.layers.Add()([x, x3])
        # # out
        x = tf.image.resize(x,
                            size=self.shape,
                            method=self.resize_method,
                            preserve_aspect_ratio=False,
                            antialias=False,
                            name=None)
        x = tf.keras.layers.Activation('sigmoid')(x)
        # x = tf.nn.softmax(x, axis=-1)
        return x
